import pybuda._C
from typing import ClassVar, Dict, Iterator, List, Optional, Tuple, Union, overload

C: UBlockOrder
Concatenate: RuntimeTensorTransformType
ConstantInput: RuntimeTensorTransformType
EmbeddingIndex: RuntimeTensorTransformType
NoTransform: RuntimeTensorTransformType
Prestride: RuntimeTensorTransformType
R: UBlockOrder
ReinterpretShape: RuntimeTensorTransformType
Unpad: RuntimeTensorTransformType
kBudaNaryTM: NodeType
kBudaOp: NodeType
kInput: NodeType
kOutput: NodeType
kPyOp: NodeType
kQueue: NodeType

class Graph:
    def __init__(self, arg0: str) -> None: ...
    def clone(self) -> Graph: ...
    def enable_training(self) -> bool: ...
    def get_constant_input_runtime_tensor_transform_constants(self) -> List[Tuple[str, object]]: ...
    def get_constant_names(self) -> List[str]: ...
    def get_constant_nodes(self, *args, **kwargs): ...
    def get_fused_ops(self) -> List[Tuple[int, List[List[str]]]]: ...
    def get_input_runtime_tensor_transforms(self, *args, **kwargs): ...
    def get_microbatch(self) -> int: ...
    def get_name(self) -> str: ...
    def get_node_id(self, arg0: str) -> int: ...
    def get_node_name(self, arg0: int) -> str: ...
    def get_ordered_constant_tile_dims(self) -> List[List[int]]: ...
    def get_ordered_input_gradient_names(self) -> List[str]: ...
    def get_ordered_input_names(self) -> List[str]: ...
    def get_ordered_input_requires_grad(self) -> List[bool]: ...
    def get_ordered_input_shapes(self) -> List[List[int]]: ...
    def get_ordered_input_subgraph_indices(self) -> List[int]: ...
    def get_ordered_input_tile_dims(self) -> List[List[int]]: ...
    def get_ordered_intermediate_names(self) -> List[str]: ...
    def get_ordered_intermediate_shapes(self) -> List[List[int]]: ...
    def get_ordered_output_gradient_names(self) -> List[str]: ...
    def get_ordered_output_names(self) -> List[str]: ...
    def get_ordered_output_requires_grad(self) -> List[bool]: ...
    def get_ordered_output_shapes(self) -> List[List[int]]: ...
    def get_ordered_output_subgraph_indices(self) -> List[int]: ...
    def get_ordered_parameter_tile_dims(self) -> List[List[int]]: ...
    def get_ordered_target_names(self) -> List[str]: ...
    def get_ordered_target_shapes(self) -> List[List[int]]: ...
    def get_ordered_target_subgraph_indices(self) -> List[int]: ...
    def get_output_runtime_tensor_transforms(self, *args, **kwargs): ...
    def get_parameter_nodes(self, *args, **kwargs): ...
    def get_subgraph_id_for_node(self, arg0: int) -> int: ...
    def get_tile_broadcast_dims_for_bw_input(self, arg0: int) -> List[int]: ...
    def get_tile_broadcast_dims_for_input(self, arg0: int) -> List[int]: ...
    def get_tile_broadcast_dims_for_target(self, arg0: int) -> List[int]: ...
    def has_node_with_id(self, arg0: int) -> bool: ...
    def nodes(self) -> List[str]: ...
    def output_node_redirected(self) -> bool: ...
    def register_module_inputs(self, arg0: List[int]) -> None: ...
    def register_module_outputs(self, arg0: List[int], arg1: List[bool]) -> None: ...
    def register_module_targets(self, arg0: List[int]) -> None: ...
    def set_enable_training(self, arg0: bool) -> None: ...
    def set_microbatch(self, arg0: int) -> None: ...

class InputNode:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def id(self) -> int: ...
    @property
    def name(self) -> str: ...
    @property
    def node_type(self): ...
    @property
    def output_df(self) -> pybuda._C.DataFormat: ...
    @property
    def shape(self) -> Shape: ...

class Node:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def id(self) -> int: ...
    @property
    def name(self) -> str: ...
    @property
    def node_type(self): ...
    @property
    def output_df(self) -> pybuda._C.DataFormat: ...
    @property
    def shape(self) -> Shape: ...

class NodeContext:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def id(self) -> int: ...
    @property
    def name(self) -> str: ...
    @property
    def node_type(self): ...
    @property
    def output_df(self) -> pybuda._C.DataFormat: ...
    @property
    def shape(self) -> Shape: ...
    @property
    def unbroadcast_shape(self) -> Shape: ...

class NodeType:
    __members__: ClassVar[dict] = ...  # read-only
    __entries: ClassVar[dict] = ...
    kBudaNaryTM: ClassVar[NodeType] = ...
    kBudaOp: ClassVar[NodeType] = ...
    kInput: ClassVar[NodeType] = ...
    kOutput: ClassVar[NodeType] = ...
    kPyOp: ClassVar[NodeType] = ...
    kQueue: ClassVar[NodeType] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class OpType:
    def __init__(self, op: str, attr=..., named_attrs=...) -> None: ...
    def set_buda_attr(self, arg0: str, arg1) -> None: ...
    def __getattr__(self, name): ...
    def __setattr__(self, arg0: str, arg1) -> None: ...
    @property
    def attr(self): ...
    @property
    def buda_attrs(self): ...
    @property
    def named_attrs(self): ...
    @property
    def op(self) -> str: ...

class RuntimeTensorTransform:
    concat_dim: int
    concat_group: int
    concat_index: int
    kernel_height: int
    kernel_width: int
    original_shape: Shape
    reinterpreted_shape: Shape
    stride_height: int
    stride_width: int
    type: RuntimeTensorTransformType
    unpadded_shape: Shape
    def __init__(self) -> None: ...
    def EmbeddingIndex(self) -> RuntimeTensorTransform: ...
    def from_json(self) -> RuntimeTensorTransform: ...
    def to_json(self) -> json: ...

class RuntimeTensorTransformType:
    __members__: ClassVar[dict] = ...  # read-only
    Concatenate: ClassVar[RuntimeTensorTransformType] = ...
    ConstantInput: ClassVar[RuntimeTensorTransformType] = ...
    EmbeddingIndex: ClassVar[RuntimeTensorTransformType] = ...
    NoTransform: ClassVar[RuntimeTensorTransformType] = ...
    Prestride: ClassVar[RuntimeTensorTransformType] = ...
    ReinterpretShape: ClassVar[RuntimeTensorTransformType] = ...
    Unpad: ClassVar[RuntimeTensorTransformType] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class Shape:
    class Type:
        __members__: ClassVar[dict] = ...  # read-only
        BUDA: ClassVar[Shape.Type] = ...
        FREE: ClassVar[Shape.Type] = ...
        __entries: ClassVar[dict] = ...
        def __init__(self, value: int) -> None: ...
        def __eq__(self, other: object) -> bool: ...
        def __hash__(self) -> int: ...
        def __index__(self) -> int: ...
        def __int__(self) -> int: ...
        def __ne__(self, other: object) -> bool: ...
        @property
        def name(self) -> str: ...
        @property
        def value(self) -> int: ...
    BUDA: ClassVar[Shape.Type] = ...
    FREE: ClassVar[Shape.Type] = ...
    def __init__(self, *args, **kwargs) -> None: ...
    def as_list(self) -> List[int]: ...
    @classmethod
    def create(cls, values: List[int]) -> Shape: ...
    @classmethod
    def create_buda(cls, arg0: List[int], arg1: int, arg2: int) -> Shape: ...
    @classmethod
    def create_with_type_from_other(cls, other: Shape, values: List[int]) -> Shape: ...
    @classmethod
    def from_json(cls, arg0: json) -> Shape: ...
    def get_tile_dim(self, *args, **kwargs): ...
    def get_tile_height(self) -> int: ...
    def get_tile_width(self) -> int: ...
    def len(self) -> int: ...
    def to_json(self) -> json: ...
    def __eq__(self, arg0: Shape) -> bool: ...
    def __getitem__(self, arg0: int) -> int: ...
    def __iter__(self) -> Iterator: ...
    def __len__(self) -> int: ...
    def __setitem__(self, arg0: int, arg1: int) -> None: ...
    @property
    def c(self) -> int: ...
    @property
    def ct(self) -> int: ...
    @property
    def r(self) -> int: ...
    @property
    def rt(self) -> int: ...
    @property
    def v(self) -> int: ...
    @property
    def w(self) -> int: ...
    @property
    def z(self) -> int: ...

class UBlockOrder:
    __members__: ClassVar[dict] = ...  # read-only
    C: ClassVar[UBlockOrder] = ...
    R: ClassVar[UBlockOrder] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

def add_partial_datacopy_edge(arg0: Graph, arg1: int, arg2: int, arg3: int, arg4: int) -> None: ...
def add_subgraph_io_link_edge(arg0: Graph, arg1: int, arg2: int, arg3: int, arg4: int) -> None: ...
def create_activation_input(arg0: Graph, arg1: str, arg2: List[int], arg3: bool, arg4: pybuda._C.DataFormat, arg5: int) -> int: ...
@overload
def create_constant_input(arg0: Graph, arg1: str, arg2: float, arg3: pybuda._C.DataFormat, arg4: int) -> int: ...
@overload
def create_constant_input(arg0: Graph, arg1: str, arg2: object, arg3: List[int], arg4: pybuda._C.DataFormat, arg5: int) -> int: ...
def create_control_edge(arg0: Graph, arg1: int, arg2: int, arg3: int, arg4: int) -> None: ...
def create_data_edge(arg0: Graph, arg1: int, arg2: int, arg3: int, arg4: int, arg5: List[tuple]) -> None: ...
def create_op_node(arg0: Graph, arg1: str, arg2: OpType, arg3: List[int], arg4: pybuda._C.DataFormat, arg5: int, arg6: Dict[str, Union[bool, int, str]]) -> int: ...
def create_output(arg0: Graph, arg1: str, arg2: List[int], arg3: pybuda._C.DataFormat, arg4: bool, arg5: int) -> int: ...
def create_parameter_input(arg0: Graph, arg1: str, arg2: List[int], arg3: bool, arg4: pybuda._C.DataFormat, arg5: int) -> int: ...
def create_target_input(arg0: Graph, arg1: str, arg2: List[int], arg3: bool, arg4: pybuda._C.DataFormat, arg5: int) -> int: ...
def eval(graph: Graph, inputs: List[object], parameters: Dict[str, object], tt_device: object, relative_atol: float, pcc: float, intermediate_golden_tensors: Dict[int, object] = ..., losses: List[object] = ..., targets: List[object] = ..., balancer_solution=..., dump_tensors_path: str = ..., allow_modified_shapes: bool = ...) -> Tuple[List[object], Dict[str, object], List[object], Dict[str, object]]: ...
def get_constant_input_value(arg0: Node, arg1: bool) -> object: ...
def get_intermediate_tensors(graph: Graph, inputs: List[object], parameters: Dict[str, object], tt_device: object, relative_atol: float, pcc: float, intermediate_golden_tensors: Dict[int, object] = ..., losses: List[object] = ..., targets: List[object] = ..., balancer_solution=..., dump_tensors_path: str = ..., allow_modified_shapes: bool = ...) -> Dict[str, object]: ...
def get_optimizer_param_info(arg0: Graph, arg1: str) -> List[Tuple[InputNode, str]]: ...
def get_shape_for_node(arg0: Graph, arg1: str) -> List[int]: ...
def record_consteval_operations(arg0: Graph) -> Dict[str, Optional[json]]: ...
def remove_node(arg0: Graph, arg1: int) -> None: ...
